module Magicloud.JobGroup where

import Magicloud.Jobs
import qualified Data.Map as M
import qualified Magicloud.Map as Map
import qualified Control.Concurrent.MVar.ReadOnly as ROMVar
import Control.Concurrent
import Control.Exception

data StartKill k a = StartKill { toStart :: JobArgs k a
                               , toKill :: [k] }

data JobGroupInfo k e a = JobGroupInfo { jobGroupId :: ThreadId
                                       , jobs :: ROMVar.ReadOnlyMVar (Jobs k e a) }

start :: (Ord k, Exception e) => IO (StartKill k a) -> (a -> IO b) -> (k -> Either e b -> IO ()) -> IO (JobGroupInfo k e b)
start sk worker cb = mask_ $ do
  jobs_ <- newMVar M.empty
  tId <- forkIOWithUnmask $ \unmask ->
    unmask $ start_ sk worker cb jobs_
  return $ JobGroupInfo tId $ ROMVar.toReadOnlyMVar jobs_
  where
    start_ :: (Ord k, Exception e) => IO (StartKill k a) -> (a -> IO b) -> (k -> Either e b -> IO ()) -> MVar (Jobs k e b) -> IO ()
    start_ sk_ worker_ cb_ jobs_ = mask $ \restore -> do
      runningJobs <- (readMVar jobs_ >>=
        Map.filterWithKeyM (\k (JobInfo _ st) -> do
                             st_ <- ROMVar.tryTakeMVar st
                             case st_ of
                               Nothing -> return True
                               Just st__ -> do
                                 _ <- forkIOWithUnmask $ \unmask ->
                                   unmask $ cb_ k st__
                                 return False
                           ))
      StartKill ss ks <- restore $ sk_
      mapM_ (\k ->
              case (M.lookup k runningJobs) of
                Just (JobInfo tId _) -> killThread tId
                Nothing -> return ()
            ) ks
      newJobs <- restore $ Magicloud.Jobs.start ss worker_
      _ <- swapMVar jobs_ (runningJobs `M.union` newJobs)
      start_ sk_ worker_ cb_ jobs_

kill :: JobGroupInfo k e a -> IO ()
kill (JobGroupInfo jgId js) = do
  jobs_ <- ROMVar.takeMVar js
  mapM_ (\(JobInfo jId _) -> killThread jId) $ M.elems jobs_
  killThread jgId
